#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include "mpi.h"

#define Max(a,b) ((a)>(b)?(a):(b))
#define PLOT_FILE "plots"   

double pi = 3.14159265358979323846;
int size, rank;
double **A;
int N = 512;

void relax(int localSize, double w, double *eps, int i0 );
void init(int localSize);

void update_bc( int localSize,int below, int above );
double getW(double w, int iter);
int write_file(int localSize);

int main(int argc, char *argv[]) {
	double time0, time1;
	double w = 0;
	double maxeps = 0.01;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &size); 
    MPI_Comm_rank(MPI_COMM_WORLD, &rank); 

	if( rank == 0 ){		
		if (argc > 1) 
			N = atoi(argv[1]); 
		
		if (argc > 2)
			maxeps = atof(argv[2]); 
			
		if( argc > 3 ){
			fprintf(stdout, "Too much arguments");
			return 1;
		}
		
		time0 = MPI_Wtime();
        fprintf(stdout, "size=%d, N = %d, maxeps = %.3f\n", size, N, maxeps);
	}

	MPI_Bcast( &N     , 1, MPI_INT   , 0, MPI_COMM_WORLD );
	MPI_Bcast( &maxeps, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD );

	int localSize = N / size;

    init(localSize+2);
     
    int below, above;
    if (rank == 0) {
         below = MPI_PROC_NULL;
         above = rank + 1;
    }
    else if (rank == size - 1) {
         below = rank - 1;
         above = MPI_PROC_NULL;
    }
    else {
         below = rank - 1;
         above = rank + 1;
    }

    double eps, eps1, eps2, geps = 1.0;
    int iter = 0;
    while( iter < 100000 ) {
	
		eps1 = 0;
        w = getW(w, iter);
        relax(localSize + 2, w, &eps1, 1);
        update_bc(localSize + 2, below, above);

        iter++;

		eps2 = 0;
        w = getW(w, iter);
        relax(localSize + 2, w, &eps2, 0);
        update_bc(localSize + 2, below, above);

		eps = eps1 + eps2;

		MPI_Allreduce(&eps, &geps, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); 

        if (geps < maxeps)
			break;
	}
    MPI_Barrier(MPI_COMM_WORLD);

	if( rank == 0 ){
		time1 = MPI_Wtime();
        fprintf(stdout, "Time: = %gs\t\n", time1 - time0 );
	}

    write_file(localSize);

    MPI_Barrier(MPI_COMM_WORLD);

    for( int i = 0; i < localSize + 2; i++ )
		free(A[i]);
    free(A); 

	MPI_Finalize();
	return 0;
}

void init(int localSize) {
    int N2 = N+2;

    A = (double**)malloc(localSize*sizeof(double*));
    for (int i = 0; i < localSize; i++) {
        A[i] = (double*)malloc(N2*sizeof(double));

        for (int j = 0; j < N2; j++) {
            A[i][j] = 0;
        }
    }

    if (rank == 0) {
        for (int j = 0; j < N2; j++) {
            A[0][j] = sin(pi * ( (double)j / N2 ) * exp( -( (double)j / N2 ) ));
        }
    }

    if (rank == size - 1) {
        for (int j = 0; j < N2; j++) {
            A[localSize - 1][j] = sin(pi * ( (double)j / N2 ) );
        }
    }

}


void relax(int localSize, double w, double *eps, int i0) {
    int N2 = N+2;
	double temp, Anew;

    for (int i = 1; i < localSize - 1; i++) {
        for (int j =  1 + (i + i0) % 2; j < N2 - 1; j += 2) {
            temp = (A[i - 1][j] + A[i + 1][j] + A[i][j - 1] + A[i][j + 1]) / 4.0;
            Anew = (1.0 - w) * A[i][j] + w * temp;
            *eps += fabs(Anew - A[i][j]);
            A[i][j] = Anew;
        }
    }
}

void update_bc(int localSize, int below, int above) {

    MPI_Sendrecv( A[localSize - 2] + 1, N, MPI_DOUBLE, above, 0,
                  A[0] 			   + 1, N, MPI_DOUBLE, below, 0,
				  MPI_COMM_WORLD, MPI_STATUS_IGNORE 
				 );

    MPI_Sendrecv( A[1]             + 1, N, MPI_DOUBLE, below, 1,
                  A[localSize - 1] + 1, N, MPI_DOUBLE, above, 1,
				  MPI_COMM_WORLD, MPI_STATUS_IGNORE 
				 );
}

double getW(double w, int iter){
    double p = pi / (2.0 * (N + 1));
    double ro = 1-p*p;

    if (iter == 0) return 0.0;
    if (iter == 1) return 1 / (1 - ro * ro / 2);
    return 1 / (1 - w * ro * ro / 4);
}

int write_file(int localSize) {
    char filename[50], file[53];
    FILE *fd;
    int N2 = N+2;

    sprintf(filename, "%s", PLOT_FILE);

    if (rank == 0) {
        fd = fopen(filename, "w");
        fprintf(fd, "%5d %5d %5d\n", N2, N2, size);
        fclose(fd);
    }

    sprintf(file, "%s.%d.%d.%d", filename, N, size, rank);
    fd = fopen(file, "w");

    int ij = 0;
    for (int j = 0; j < N2; j++) {
        for (int i = 0; i < localSize + 2; i++) {
            fprintf(fd, "%11.4f ", A[i][j]);

            if ((ij + 1) % 6 == 0)
                fprintf(fd, "\n");

            ij++;
        }
    }

    fprintf(fd, "\n");
    fclose(fd);
    return 0;
}
